#!/bin/bash
if [ $# -lt 1 ]; then
    showhelp
    exit 1
fi

function showhelp {
    echo 'usage : orcas option *.inp'
    echo 'e.g.  : orcas -n 32 -e node05 -E *.gjf'
    echo 'option: -[n number of cores]'
    echo '        -[m total memory in GB] '
    echo '        -[M total memory in percent, 80 means 80% of total memory] '
    echo '        -[e nodename to exclude. e.g. -e node01,node02]'
    echo '        -[q queue name to submit]'
    echo '        -[E exclusive mode]'
    echo '        -[p pbs or slurm]'
    echo '        -[t tmpdir]'

}

while getopts ":n:l:p:q:t:c:e:Em:M:" opt
do
    case $opt in
        n) NPROC=$OPTARG;;
        p) SCHED=$OPTARG;;
        c) CHECK=$OPTARG;;
        q) QUEUE=$OPTARG;;
        e) EXCLUDE=$OPTARG;;
        E) EXCLUSIVE=on;;
        l) NODE=$OPTARG;;
        t) TEMPDIR=$OPTARG;;
        m) MEM=$OPTARG;;
        M) MPERCENT=$OPTARG;;
        ?) showhelp
        exit 1;;
    esac
done
shift $(($OPTIND - 1))

while [[ $# -ge 1 ]]
do
dos2unix $1
RUN_FLAG=1
if [ -z "$NPROC" ]; then
    XNPROC=`awk /nprocs/'{if($2=="nprocs") print $3;if($1=="nprocs") print $2}' $1`
else
    sed -i "s/%pal nprocs.*/%pal nprocs $NPROC/" $1
    XNPROC=$NPROC
fi
if [ -n "$MEM" ]; then
    MEMMB=`echo "$MEM * 1024 / $XNPROC" | bc`
    sed -i "s/%MaxCore.*/%MaxCore $MEMMB/i" $1
fi
if [ -z "$SCHED" ];then
    SCHED=slurm
fi

if [ -z "$TEMPDIR" ];then
    TEMPDIR='/data'
fi
if [[ "$CHECK" == "2" ]];then
    NORM_FLAG=`grep  NORMALLY ${1/.inp/}.out`
    if [ -n "$NORM_FLAG" ];then
        RUN_FLAG=0
    else
        echo "${1/.inp/}.out not terminated normally. Submitting..."
    fi
fi
if [[ "$CHECK" == "1" && -f ${1/.inp/}.out ]];then
    if [ -f ${1/.inp/}.out ];then
        RUN_FLAG=0
    else
        echo "${1/.gjf/}.log not exist. Submitting..."
    fi 
fi

xyzfile=`grep -i xyzfile $1 | awk '{print $NF}'`

#ORCA=`which orca`


# check job manage system
PBS_FLAG=`which qsub 2> null`
SLURM_FLAG=`which sbatch 2> null`
if [ -n "$SLURM_FLAG" ] && [ "$SCHED" == "slurm" ];then
    if [[ "$RUN_FLAG" == "1" ]]; then
        JOB_NAME=${1/.inp}
        NAME=sub_${1/.inp/}.sbatch
        echo '#!/bin/bash' > $NAME
        echo "#SBATCH -J $JOB_NAME" >> $NAME
        echo "#SBATCH -N 1" >> $NAME
        echo "#SBATCH -n $XNPROC" >> $NAME
        echo "#SBATCH --cpus-per-task 1" >> $NAME
        if [ -n "$QUEUE" ];then
            echo "#SBATCH --partition=$QUEUE" >> $NAME
        fi
        if [ -n "$NODE" ];then
            echo "#SBATCH -w $NODE" >> $NAME
        fi
        if [ -n "$EXCLUDE" ];then
            echo "#SBATCH --exclude $EXCLUDE" >> $NAME
        fi
        if [[ "$EXCLUSIVE" == "on" ]];then
            echo "#SBATCH --exclusive" >> $NAME
        fi
        echo "source ~/.bashrc" >> $NAME
        echo "tdir=${TEMPDIR}/${1/.inp/}" >> $NAME
        echo "mkdir \$tdir" >> $NAME
        echo "cp $1 \$tdir" >> $NAME
    
        if [ -n "$xyzfile" ];then
            echo "cp $xyzfile \$tdir" >> $NAME
        fi
        echo "cp ${1/.inp/}.gbw \$tdir" >> $NAME
        echo "cp ${1/.inp/}.*nat \$tdir" >> $NAME
        echo "cp ${1/.inp/}.pot \$tdir" >> $NAME
        echo "cd \$tdir" >> $NAME
        echo "start=\$(date +%s.%N)" >> $NAME
        if [ -n "$MPERCENT" ]; then
            echo "MEM=\$(free -h | awk -v np=$XNPROC -v mp=$MPERCENT '/Mem/{gsub(/Gi/,\"\",\$2);print \$2*1024/np*mp/100}')" >> $NAME
            echo "sed -i \"s/%MaxCore.*/%MaxCore \$MEM/gi\" $1" >> $NAME
        fi
        echo "/home/software/program/orca_5_0_3_linux_x86-64_shared_openmpi411/orca $1 2>&1 | tee \$SLURM_SUBMIT_DIR/${1/.inp/}.out" >> $NAME
        echo "end=\$(date +%s.%N)" >> $NAME
        echo "cp ${1/.inp/}.* \$SLURM_SUBMIT_DIR" >> $NAME
        echo "cd \$SLURM_SUBMIT_DIR" >> $NAME
        echo "rm -rf \$tdir" >> $NAME
        echo "optcyc=\$(grep -c 'OPTIMIZATION CYCLE' ${1/.inp/}.out)" >> $NAME
        echo 'runtime=$(printf "%.2f\n" $(echo "scale=2;$end - $start" | bc ))' >> $NAME
        echo 'cyctime=$(echo "scale=2;$runtime / $optcyc " | bc )' >> $NAME
        echo 'echo "RUN STATS: $runtime seconds; $optcyc cycles; $cyctime second/cycle" > ${1/.inp/}.runstats' >> $NAME
        sbatch $NAME
    fi 
elif [ -n "$PBS_FLAG" ] && [ "$SCHED" == "pbs" ];then
    if [[ "$RUN_FLAG" == "1" ]]; then
        JOB_NAME=${1/.inp}
        NAME=sub_${1/.inp/}.pbs
if [ -z "$NODE" ];then
    NODE=1
fi
cat << ENDOF > $NAME
#PBS -S /bin/bash
#PBS -l nodes=$NODE:ppn=44
#PBS -N $JOB_NAME
#PBS -j oe
cd \$PBS_O_WORKDIR
source ~/.bashrc
hn=\$(hostname)
tdir=/data\${hn/node/}/${1/.inp}
mkdir \$tdir
cp $1 \$tdir
cp $xyzfile \$tdir
cp ${1/.inp/}.gbw \$tdir
cp ${1/.inp/}.pot \$tdir
cd \$tdir
start=\$(date +%s.%N)
$ORCA $1 2>&1 | tee \$PBS_O_WORKDIR/${1/.inp/}.out
end=\$(date +%s.%N)
cp ${1/.inp/}.* \$PBS_O_WORKDIR
cd \$PBS_O_WORKDIR
rm -rf \$tdir
optcyc=\$(grep -c 'OPTIMIZATION CYCLE' ${1/.inp/}.out)
runtime=\$(printf "%.2f\n" \$(echo "scale=2;\$end - \$start" | bc ))
cyctime=\$(echo "scale=2;\$runtime / \$optcyc " | bc )
echo "RUN STATS: \$runtime seconds; \$optcyc cycles; \$cyctime second/cycle" > ${1/.inp/}.orunstats
ENDOF
qsub $NAME
fi
fi
shift 1
done 
